import yaml, sys, re, argparse, json
from collections import defaultdict
from typing import Dict, List, Tuple, Optional
ROBOT_GROUPS = [('aloha', 'a'), ('franka', 'f'), ('ur5', 'u'), ('piper', 'p')]

def parse_status(status_str: str) -> Tuple[float, int, int]:
    assert status_str and status_str != 'null', f"Invalid status_str: {status_str}"
    match = re.search(r'exist:(\d+)/(\d+)\s*\(([\d.]+)%\)', status_str)
    if not match: return -1, -1, -1
    success, total, percentage = int(match.group(1)), int(match.group(2)), float(match.group(3))
    return percentage, success, total

def extract_group_from_source(source: str) -> Optional[str]:
    assert source, f"Source is empty: {source}"
    match = re.search(r'-g\s+([a-zA-Z0-9_-]+)', source)
    # print(f"Source: {source}, Match: {match.group(1)}")
    return match.group(1) if match else None

def extract_dim_from_note(note: str) -> int:
    assert note, f"Note is empty: {note}"
    match = re.search(r'\((\d+)\)', note)
    return int(match.group(1)) if match else -1

def is_groupname_fine(groups: List[str]) -> bool:
    assert len(groups) in [3, 4], f"Groups length mismatch: {len(groups)}"
    assert all(groups), f"Empty groups found: {groups}"
    first_chars = [g[0] for g in groups]
    valid_chars = {'a', 'f', 'u', 'p'}
    assert set(first_chars).issubset(valid_chars), f"Invalid first chars: {first_chars}"
    suffixes = [g[1:] for g in groups]
    if len(set(suffixes)) == 1: return True
    base_suffix = suffixes[0]
    return all(base_suffix in suffix or suffix in base_suffix for suffix in suffixes[1:])

def get_eval_data(yaml_path: str) -> Tuple[List[Dict], List[str]]:
    with open(yaml_path, 'r') as file: data = yaml.safe_load(file)
    eval_section = data.get('paths', {}).get('eval', {})
    assert eval_section, f"No eval section found in {yaml_path}"
    note_groups = defaultdict(list)
    for idx, item_data in eval_section.items():
        note = item_data.get('note')
        if note: note_groups[note].append((idx, item_data))
    neglect_list, groupname_fine_notes = [], set()
    for note, items in note_groups.items():
        if len(items) in [3, 4]:
            groups = [extract_group_from_source(item[1].get('source')) for item in items]
            if is_groupname_fine(groups):
                groupname_fine_notes.add(note)
            else:
                # print(f"Groupname is not fine: {note}")
                # print(f"items: {json.dumps(items, indent=4)}")
                neglect_list.append(note)
        else: 
            # print(f"Items length is not 3 or 4: {note} {len(items)}")
            # print(f"items: {json.dumps(items, indent=4)}")
            neglect_list.append(note)
    eval_data_list = []
    for note, items in note_groups.items():
        if note in groupname_fine_notes:
            indices = [item[0] for item in items]
            robot_data = {robot_name: {'val': 0.0, 'success': 0, 'total': 0} for robot_name, _ in ROBOT_GROUPS}
            for idx, item_data in items:
                group = extract_group_from_source(item_data.get('source'))
                status = item_data.get('status', '')
                percentage, success, total = parse_status(status)
                for robot_name, prefix in ROBOT_GROUPS:
                    if group and group.startswith(prefix):
                        robot_data[robot_name].update({'val': percentage, 'success': success, 'total': total})
                        break
            valid_scores = [robot_data[robot]['val'] for robot, _ in ROBOT_GROUPS if robot_data[robot]['val']]
            avg_score = sum(valid_scores) / len(valid_scores) if valid_scores else 0.0
            dim = extract_dim_from_note(note)
            result = {'indices': indices, 'note': note.replace(f"({dim})"," "), 'avg': avg_score, 'dim': dim}
            for robot_name, _ in ROBOT_GROUPS:
                result.update({robot_name: robot_data[robot_name]['val'], f'{robot_name}_success': robot_data[robot_name]['success'], f'{robot_name}_total': robot_data[robot_name]['total']})
            eval_data_list.append(result)
    return eval_data_list, neglect_list

def analyze_eval_data(yaml_path: str, sort_by: str = None) -> None:
    eval_data_list, neglect_list = get_eval_data(yaml_path)
    assert eval_data_list, f"No eval data found in {yaml_path}"
    if sort_by:
        sort_fields = sort_by
        def sort_key(item):
            key_parts = []
            for field in sort_fields:
                if field == 'indexs':
                    key_parts.append(min(item['indices']))
                elif field == 'note':
                    note_first_part = item['note'].split(' ')[0] if ' ' in item['note'] else item['note']
                    key_parts.append(note_first_part)
                elif field in [robot[0] for robot in ROBOT_GROUPS] + ['avg']:
                    key_parts.append(-item[field] if item[field] else 0)
                elif field == 'dim':
                    key_parts.append(-item[field] if item[field] != -1 else 0)
                else:
                    key_parts.append(0)
            return tuple(key_parts)
        eval_data_list.sort(key=sort_key)
    max_values = {}
    for robot_name, _ in ROBOT_GROUPS:
        max_values[robot_name] = max(item[robot_name] for item in eval_data_list)
    max_values['avg'] = max(item['avg'] for item in eval_data_list)
    header_cols = ['indexs', 'note', 'dim'] + [robot[0] for robot in ROBOT_GROUPS] + ['avg']
    header_str = " ".join(f"{col:<17}" for col in header_cols)
    print(header_str)
    print("-" * len(header_str))
    for item in eval_data_list:
        index_range = f"{min(item['indices'])}-{max(item['indices'])}"
        robot_strings = {}
        for robot_name, _ in ROBOT_GROUPS:
            if item[robot_name]:
                if item[f'{robot_name}_total'] == 100:
                    robot_strings[robot_name] = f"{item[robot_name]:.1f}%"
                else:
                    robot_strings[robot_name] = f"{item[robot_name]:.1f}%({item[f'{robot_name}_success']}/{item[f'{robot_name}_total']})"
                if item[robot_name] == max_values[robot_name]:
                    robot_strings[robot_name] = f"[{robot_strings[robot_name]}]"
            else:
                robot_strings[robot_name] = ""
        avg_str = f"{item['avg']:.1f}%" if item['avg'] else ""
        if item['avg'] == max_values['avg']:
            avg_str = f"[{avg_str}]"
        dim_str = str(item['dim']) if item['dim'] != -1 else ""
        output_parts = [index_range, item['note'], dim_str]
        for robot_name, _ in ROBOT_GROUPS:
            output_parts.append(robot_strings[robot_name])
        output_parts.append(avg_str)
        print(" ".join(f"{part:<17}" for part in output_parts))
    print(f"\nNeglected notes: {', '.join(neglect_list)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Analyze eval data from YAML file')
    parser.add_argument('yaml_name', help='Name of the YAML file (without .yaml extension)')
    sort_choices = ['indexs', 'note', 'dim'] + [robot[0] for robot in ROBOT_GROUPS] + ['avg']
    parser.add_argument('-s', '--sort', nargs='+', choices=sort_choices, help='Sort by specified fields (e.g., "note avg" for multi-level sorting)')
    args = parser.parse_args()
    yaml_path = f"/home/algo/geyiheng/utilgyh/{args.yaml_name}.yaml"
    analyze_eval_data(yaml_path, args.sort)
   